
#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@文件   :control_reset_flow.py
@说明   :
@时间   :2024/04/28 12:39:59
@作者   :TaoHong
@版本   :v1.0
'''


import pandas as pd
import sumolib

import xml.etree.ElementTree as ET


def pretty_xml(elem, level=0):
    '''
    XML 文档写入文件的时候
    增加换行符
    '''
    i = "\n" + level*"\t"
    if len(elem):
        if not elem.text or not elem.text.strip():
            elem.text = i + "\t"
        if not elem.tail or not elem.tail.strip():
            elem.tail = i
        for elem in elem:
            pretty_xml(elem, level+1)
        if not elem.tail or not elem.tail.strip():
            elem.tail = i
    else:
        if level and (not elem.tail or not elem.tail.strip()):
            elem.tail = i


class resetFlow:
    
    def __init__(self):
        """
        """


    def get_VType(self):
        """
        获取不同车辆类型的行驶特征参数
        Excel文件.“VType”中，提前给定
        车辆类型，这里需要区分的是：
            客车
            货车
        """
        
        root = self.root
        
        da = pd.read_excel(self.flow_file, sheet_name="VType", header=0, index_col=0)

        self.VType = da.iloc[0, 1:]     # 不要第一列的第一行

        for car_type in range(1, da.shape[1]):  #从1开始 多少种类型的车需要自定义/添加
            
            vTypeChild = ET.SubElement(root, 'vType')

            for car_para in range(da.shape[0] ): # 多少个参数需要自定义
                vTypeChild.set(da.iloc[car_para, 0], str(da.iloc[car_para, car_type]))
        

    def get_road_toll_edges(self, OD_id, OD_flag):
        """
        OD_edges_df: 各个OD对中包含的edges
        OD_id:       OD序号 O_id/D_id
        OD_flag:     origin/destination
        
        处理一下，给不同地区出发/到达的edge一个车辆比例
        """

        road_flag = OD_flag+'Road' # 从road（其他相交的道路）出发或到达
        toll_flag = OD_flag+'Toll' # 从toll（收费站）       出发或到达

        OD_Road = self.OD_edges_df.loc[int(OD_id), road_flag]
        OD_Toll = self.OD_edges_df.loc[int(OD_id), toll_flag]

        # 处理一下，给一个比例
        if pd.isna(OD_Road): 
            OD_edge_ls = OD_Toll.split(';')
            OD_df = pd.DataFrame({'OD_edge': OD_edge_ls,
                                  'OD_rate': [1]*len(OD_edge_ls)})
            
            OD_df.loc[:,'OD_rate'] /= len(OD_edge_ls)

        elif pd.isna(OD_Toll):
            OD_edge_ls = OD_Road.split(';')
            OD_df = pd.DataFrame({'OD_edge': OD_edge_ls,
                                  'OD_rate': [1]*len(OD_edge_ls)})
            
            OD_df.loc[:,'OD_rate'] /= len(OD_edge_ls)
        
        else:
            road_rate = 0.8 # road和toll  8:2
            toll_rate = 0.2
            #! Road
            OD_edge_ls = OD_Road.split(';')
            OD_df_1 = pd.DataFrame({'OD_edge': OD_edge_ls,
                                  'OD_rate': [road_rate]*len(OD_edge_ls)})
            
            OD_df_1.loc[:,'OD_rate'] /= len(OD_edge_ls)

            #! Toll
            OD_edge_ls = OD_Toll.split(';')
            OD_df_2 = pd.DataFrame({'OD_edge': OD_edge_ls,
                                  'OD_rate': [toll_rate]*len(OD_edge_ls)})
            
            OD_df_2.loc[:,'OD_rate'] /= len(OD_edge_ls)

            #! road和toll合并
            OD_df = pd.concat([OD_df_1, OD_df_2]
                              ,ignore_index=True
                              ,axis=0
                    )

        return OD_df



    def add_Vehicle_route(self):
        """
        先加route 之后再加flow

        获取OD对之间的交通流量
        Excel文件中,“OD-Flow”中提前给定
        """
        
        root = self.root

        # 通过预定义/提前给定的OD对和edgeID对应关系表sheet“OD-Edges”
        # 获取各个OD对中包含的edges, 用作流量输入输出flow的 from 和 to
        da = pd.read_excel(self.flow_file, sheet_name="OD-Edges", header=0, index_col=0, dtype=str)
        self.OD_edges_df = da.iloc[:,-4:]

        net = sumolib.net.readNet(self.net_file, withInternal=False)

        for O_id in da.index: #* 是int类型 从1开始
            OD_flag = 'origin'
            origin_edges_df = self.get_road_toll_edges(O_id, OD_flag)
            
            for D_id in da.index: #* 是int类型 从1开始
                
                if O_id == D_id: # O和D的相同，自己走普通道路绕路吧，别来高速沾边
                    continue

                OD_flag = 'destination'
                destination_edges_df = self.get_road_toll_edges(D_id, OD_flag)

                for i in range(origin_edges_df.shape[0]): 
                    for j in range(destination_edges_df.shape[0]):
                        
                        begE = net.getEdge(origin_edges_df.loc[i,'OD_edge'])
                        endE = net.getEdge(destination_edges_df.loc[j, 'OD_edge'])

                        routeTuple = net.getShortestPath(begE, endE)[0] #! 不一定可以联通
                
                        if routeTuple != None: # 有联通的道路，最短路不为空
                            
                            #! 定义路径
                            strRoute = ' '.join([edg.getID() for edg in routeTuple])
                            
                            routeChild = ET.SubElement(root, 'route')
                            rouID = f'O{O_id}_D{D_id}_i{i+1}_j{j+1}'
                            
                            routeChild.set('id', rouID)
                            routeChild.set('edges', strRoute)
                            
                        # else: #! 看看那些路不连通
                            # rouID = f'O{O_id}_D{D_id}_i{i+1}_j{j+1}'
                            # print(rouID)


    def add_Vehicle_flow(self):
        """
        先加route 之后再加flow

        获取OD对之间的交通流量
        Excel文件中,“OD-Flow”中提前给定
        """
        
        root = self.root

        da = pd.read_excel(self.flow_file, sheet_name="OD-Flow", header=0, index_col=0)
        net = sumolib.net.readNet(self.net_file, withInternal=False)
        

        for O_id in da.index: #* 是int类型 从1开始
            OD_flag = 'origin'
            origin_edges_df = self.get_road_toll_edges(O_id, OD_flag)
            
            for D_id in da.columns: #* 是str类型 从1开始
                
                if O_id == int(D_id): # O和D的相同，自己走普通道路绕路吧，别来高速沾边
                    continue

                OD_flag = 'destination'
                destination_edges_df = self.get_road_toll_edges(D_id, OD_flag)
                
                flow_total = da.loc[O_id, D_id] # O到D的总流量
                #! 实际上是有一个误差存在的，
                #! origin_edges_df和destination_edges_df中的比例都在之前分别定义好了，加起来为1
                #! 如果不联通，就会有少出现一部分的流量，之后看看如何改进
                
                for i in range(origin_edges_df.shape[0]): 
                    for j in range(destination_edges_df.shape[0]):
                        
                        begE = net.getEdge(origin_edges_df.loc[i,'OD_edge'])
                        endE = net.getEdge(destination_edges_df.loc[j, 'OD_edge'])

                        routeTuple = net.getShortestPath(begE, endE)[0] #! 不一定可以联通
                
                        if routeTuple != None: # 有联通的道路，最短路不为空
                            
                            rouID = f'O{O_id}_D{D_id}_i{i+1}_j{j+1}'
                            
                            #! 定义流量
                            for t in range(len(self.VType)): # 每个路径上，都添加上2种车型
                                veh_type = self.VType[t]
                                flowChild = ET.SubElement(root, 'flow') 
                                    
                                flowChild.set('id', rouID + '_' + veh_type)
                                flowChild.set('type', veh_type)    # 车辆类型
                                flowChild.set('route', rouID)
                                flowChild.set('begin', '0')
                                flowChild.set('end', '3600')
                                flowChild.set('departLane','best')
                                # flowChild.set('color', "1,0,0")
 
                                flow_i_j = round(flow_total
                                                * origin_edges_df.loc[i,'OD_rate']
                                                * destination_edges_df.loc[j, 'OD_rate']
                                                ) # 对应车辆的数量、取整
                                
                                if 'trcuk' in veh_type:
                                    flow_i_j = flow_i_j*self.truck_rate_ls[O_id-1]
                                else:
                                    flow_i_j = flow_i_j*(1-self.truck_rate_ls[O_id-1])
                                
                                flowChild.set('vehsPerHour', str(flow_i_j))
                        
                        
    

    def add_emerg_vehicle(self):
        """"
        增加事故车辆
        """

        vehicleChild = ET.SubElement(self.root, 'vehicle')
        vehicleChild.set('id', self.emerg_id)
        vehicleChild.set('type', self.VType[0])    # 车辆类型
        vehicleChild.set('route', 'O1_D3_i1_j1') # 直接定义一个事故车的路径
        vehicleChild.set('depart', '0') # 进入时间
        vehicleChild.set('color', "1,0,0")
        vehicleChild.set('departLane','best')
        vehicleChild.set('departSpeed', 'max')


    def write_route_file(self):
        # 文件生成
        pretty_xml(self.root) # 增加换行

        tree = ET.ElementTree(self.root)
        tree.write(self.rou_file
                #    ,pretty_print=True
                   ,xml_declaration=True
                   ,encoding='utf-8'
                   )


    def run_reset_flow(self, root, flow_state_dict):
        """
        增加/修改流量的功能，统一调用接口
        """
        #! 获取信息

        self.root = root
        
        # 正常车辆
        self.flow_file = flow_state_dict['flow_file']
        self.net_file = flow_state_dict['net_file']
        self.rou_file = flow_state_dict['rou_file']
        self.truck_rate_ls = flow_state_dict['truck_rate_ls']
        # 事故车
        self.emerg_id = flow_state_dict['emerg_id']

        #! 更新信息
        self.get_VType()
        
        self.add_Vehicle_route()

        self.add_emerg_vehicle()

        self.add_Vehicle_flow()
               
        self.write_route_file()


if __name__ == '__main__':

    #! 修改流量 
    print("开始修改流量！")
    root = ET.Element("routes")
    
    truck_rate_ls = [  0.3, 0.2, 0.2
                     , 0.1, 0.2, 0.1
                     , 0.1, 0.1, 0.3] # 各个Origin的大车占比
    

    flow_state_dict = {}
    
    # 正常车辆
    flow_state_dict['flow_file'] = "Flow-vType.xlsx"
    flow_state_dict['net_file'] = 'osm_road.net.xml'         # 路网
    flow_state_dict['rou_file'] = 'osm.rou.xml'
    # flow_state_dict['rou_file'] = 'new_test.rou.xml'    #* 测试流量文件
    flow_state_dict['truck_rate_ls'] = truck_rate_ls
    
    # 事故车
    flow_state_dict['emerg_id'] = 'accident_veh1'


    RF = resetFlow()
    RF.run_reset_flow(root, flow_state_dict)
    
    print("结束修改流量！")
